from zhipuapi import *
import json
import textwrap


def read_novel(file_stream, limit=500):
    with open('texts/drama_template.json') as f:
        template = json.load(f)
    
    template_string = json.dumps(template, ensure_ascii=False, separators=(',', ':'))

    prompt = textwrap.dedent(f"""
    请根据小说节选，提取故事中两个最主要的角色，并且要求对角色内容做概括总结，其中角色总结包含角色本人的基本信息和社会属性之外，也要包含跟其他角色的关系。然后生成一个场景，让提取的角色都参与进来。
    最终结果以JSON文本输出，不要包含其他多余的返回提示说明，其中JSON格式如下：{template_string}
    """)

    # streamlit.runtime.uploaded_file_manager.UploadedFile to str
    obj = file_stream.read()
    if not isinstance(obj, str):
        text = file_stream.getvalue().decode('utf-8')
    else:
        text = obj

    if limit and isinstance(limit, int) and limit > 0:
        text = text[:limit]

    glm = GlmApi(model='glm-4-air', system_prompt=prompt)
    resp, succ = glm.request(text)

    if succ:
        result = extract_json_string(resp)
        return json.loads(result)
    else:
        return None


def extract_json_string(text):
    first_square_bracket_index = 0
    last_square_bracket_index = len(text) - 1
    for i in range(0, len(text)):
        if text[i] == '{':
            first_square_bracket_index = i
            break
    for j in range(len(text)-1, 0, -1):
        if text[j] == '}':
            last_square_bracket_index = j
            break
    return text[first_square_bracket_index:last_square_bracket_index+1]
